package com.qiangesoft.datasourcepro.core;

import com.alibaba.druid.pool.DruidDataSource;
import com.qiangesoft.datasourcepro.entity.BcgDataSource;
import com.qiangesoft.datasourcepro.service.IBcgDataSourceService;
import com.qiangesoft.datasourcepro.utils.SpringUtil;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.Map;

/**
 * 动态数据源
 *
 * @author qiangesoft
 * @date 2024-03-14
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    private Map<Object, Object> dynamicTargetDataSources;

    private Object dynamicDefaultTargetDataSource;

    public DynamicDataSource(DruidDataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();

        this.dynamicTargetDataSources = targetDataSources;
        this.dynamicDefaultTargetDataSource = defaultTargetDataSource;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String datasourceId = DataSourceContext.getDataSource();
        if (datasourceId != null && this.dynamicTargetDataSources.containsKey(datasourceId)) {
            return datasourceId;
        }
        return null;
    }

    /**
     * 添加数据源
     *
     * @param dataSourceId
     */
    public void addDataSource(String dataSourceId) {
        if ("default".equals(dataSourceId)) {
            return;
        }

        int initialSize = ((DruidDataSource) dynamicDefaultTargetDataSource).getInitialSize();

        BcgDataSource bcgDataSource = SpringUtil.getBean(IBcgDataSourceService.class).getById(Long.valueOf(dataSourceId));
        if (bcgDataSource == null) {
            throw new RuntimeException("数据源配置不存在");
        }

        String datasourceId = String.valueOf(bcgDataSource.getId());
        Map<Object, Object> dynamicTargetDataSources = this.dynamicTargetDataSources;
        if (!dynamicTargetDataSources.containsKey(datasourceId)) {
            DruidDataSource dataSourceInstance = DataSourceBuilder.create()
                    .driverClassName(bcgDataSource.getDriverClassName())
                    .url(bcgDataSource.getUrl())
                    .username(bcgDataSource.getUsername())
                    .password(bcgDataSource.getPassword())
                    .type(DruidDataSource.class)
                    .build();
            dynamicTargetDataSources.put(datasourceId, dataSourceInstance);
            // 关键一步：将TargetDataSources中的连接信息放入resolvedDataSources管理
            super.afterPropertiesSet();
        }
    }
}